Refactor GroupedLinear quantization dispatch - #3326
Open
negvet wants to merge 3 commits into
Open
Conversation
Signed-off-by: Evgeny <etsykunov@nvidia.com>
Contributor
Greptile SummaryThe PR extracts GroupedLinear split-quantization logic into a dedicated module and consolidates split, quantization, and bias-gradient handling behind one dispatch API.
Confidence Score: 5/5The PR appears safe to merge because no blocking failure remains. No blocking failure remains. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
GL["GroupedLinear forward/backward"] --> SQ["_split_quantize"]
SQ --> P{"Quantized output?"}
P -->|No| Plain["Cast and torch.split"]
P -->|Yes, native bgrad| BGrad["Per-split bgrad_quantize"]
P -->|Yes, DebugQuantizer| Debug["Debug multi-tensor quantize"]
P -->|Yes, HybridQuantizer| Hybrid["Rowwise/columnwise split quantization"]
P -->|Yes, native quantizer| Native["Bulk split_quantize"]
P -->|Yes, custom/identity| Fallback["Python per-split fallback"]
Plain --> Out["Split outputs and optional dbiases"]
BGrad --> Out
Debug --> Out
Hybrid --> Out
Native --> Out
Fallback --> Out
Reviews (2): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile |
Collaborator
Author
|
/te-ci pytorch L1 |
zhongbozhu
reviewed
Aug 7, 2026
| # | ||
| # See LICENSE for license information. | ||
|
|
||
| """Grouped split-quantization helpers used by :mod:`GroupedLinear`.""" |
Collaborator
There was a problem hiding this comment.
Can we rename this file to something like split quantization? It's not really a grouped quantization
zhongbozhu
reviewed
Aug 7, 2026
| # tensors (like scales), but bulk allocation shares storage across all tensors, | ||
| # so if scales can't be offloaded, nothing in the group can be offloaded. | ||
| inputmats = _split_quantize( | ||
| inputmats, _ = _grouped_quantization._split_quantize( |
Collaborator
There was a problem hiding this comment.
Same comment for renaming
zhongbozhu
reviewed
Aug 7, 2026
| else: | ||
| input_quantizer.set_usage(rowwise=False, columnwise=True) | ||
| inputmats = _split_quantize( | ||
| inputmats, _ = _grouped_quantization._split_quantize( |
zhongbozhu
reviewed
Aug 7, 2026
zhongbozhu
left a comment
zhongbozhu
left a comment
Collaborator
There was a problem hiding this comment.
overall LGTM, just have some questions
| ) | ||
| return outputs, dbiases | ||
|
|
||
| concrete_quantizers = cast(Sequence[Quantizer], quantizers) |
Collaborator
There was a problem hiding this comment.
why do we need this cast(Sequence[Quantizer], quantizers)? feels like some extra CPU overhead?
Collaborator
Author
There was a problem hiding this comment.
This is only for the type checker, like assume there are no None elements. At runtime, it returns the same object unchanged, it does not iterate through.
Collaborator
|
/te-ci pytorch |
Signed-off-by: Evgeny <etsykunov@nvidia.com>
for more information, see https://pre-commit.ci
Collaborator
Author
|
/te-ci pytorch |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Move grouped quantization helpers into a dedicated module and replace separate split APIs with one. Keeps quantizer validation at construction only for CustomRecipe while skipping validation for the built-in recipes.
Followup for #2817
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: